from utils.array_utils import *


def partition_01(arr):
    # 小于等于数组第一个值的放左边 大于的放右边
    t = arr[0]
    s = -1
    cur = 0
    while cur < len(arr):
        if arr[cur] <= t:
            swap(arr, cur, s + 1)
            s += 1
        cur += 1


def partition_02(arr, l, r):
    # 小于第一个值的放左边, 大于的放右边, 等于的放中间
    # 返回等于区的左右下标
    t = arr[l]
    s = l - 1
    b = r + 1
    cur = l
    while cur < b:
        if arr[cur] < t:
            swap(arr, cur, s + 1)
            s += 1
            cur += 1
        elif arr[cur] > t:
            swap(arr, cur, b - 1)
            b -= 1
        elif arr[cur] == t:
            cur += 1
    # print('arr: ', arr)
    # return arr[:s + 1], arr[b:]
    return s + 1, b - 1


def partition_03(arr, l, r):
    """调整数组时保留最后一个值不动, 等都排列好后
     将最后一位和大于区最后一个值交换 返回等于区域的左右索引
     返回等于区的左右下标"""
    t = arr[r]
    cur = l
    s = l - 1
    b = r
    while cur < b:
        if arr[cur] < t:
            swap(arr, cur, s + 1)
            s += 1
            cur += 1
        elif arr[cur] > t:
            swap(arr, cur, b - 1)
            b -= 1
        elif arr[cur] == t:
            cur += 1
    swap(arr, b, r)
    return s + 1, b


"""
2022-7-14再练习荷兰国旗
"""


def partition_04(arr, l, r):
    s = l - 1
    b = r + 1
    cur = l
    t = arr[l]
    while cur < b:
        if arr[cur] < t:
            swap(arr, cur, s + 1)
            cur += 1
            s += 1
        elif arr[cur] > t:
            swap(arr, cur, b - 1)
            b -= 1
        else:
            cur += 1
    return s + 1, b - 1


def partition_sort_01(arr, l=0, r=None):
    if r is None:
        r = len(arr) - 1
    if l > r:
        return
    # i, j = partition_03(arr, l, r)
    i, j = partition_04(arr, l, r)
    # print(l, i, j, r, arr,arr[i])
    # print(arr[l:i], arr[j+1: r])
    partition_sort_01(arr, l, i - 1)
    partition_sort_01(arr, j + 1, r)


"""
2022-7-14再练习
"""


def partition_sort_02(arr, l=0, r=None):
    if r is None:
        r = len(arr) - 1
    if l >= r:
        return
    m1, m2 = partition_04(arr, l, r)
    # print(arr,m1,m2)
    partition_sort_02(arr, l, m1)
    partition_sort_02(arr, m2 + 1, r)


def swap(l, a, b):
    l[a], l[b] = l[b], l[a]


if __name__ == '__main__':
    ll = [5, 20, 6, 4, 6, 9, 1, 0, 5, 3, 6]
    # ll = [5, 20, 6, 4, 6, 9]
    print(partition_04(ll, 0, len(ll) - 1))
    print(partition_03(ll, 0, len(ll) - 1))
    # print(ll)
    # print('p3', partition_03(ll, 0, len(ll) - 1))
    # print(ll)
    partition_sort_02(ll, 0, len(ll) - 1)
    print(ll)
    print(check_all(partition_sort_02, 100))
    # for i in range(20):
    #     ll = sjlb()
    #     print(ll[0],partition_02(ll,0,len(ll)-1),ll)
